#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
验证 knowledge_base 集合的向量检索
- 随机/指定问题执行语义检索
- 打印命中片段（title/source/snippet/score）

运行：
  python knowledge_base/validate_knowledge_base.py "报销需要多久到账"
或：
  python knowledge_base/validate_knowledge_base.py
"""

import os
import sys
import logging
from typing import List, Dict, Any

import dotenv
from langchain_community.embeddings import OpenAIEmbeddings
from pymilvus import MilvusClient

# 加载环境变量
dotenv.load_dotenv()

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL")
OPENAI_EMBEDDING_MODEL = os.getenv("OPENAI_EMBEDDING_MODEL")
MILVUS_HOST = os.getenv("MILVUS_HOST")
MILVUS_PORT = os.getenv("MILVUS_PORT")
COLLECTION_NAME = "knowledge_base"


def search(query: str, top_k: int = 5):
    if not (OPENAI_API_KEY and OPENAI_BASE_URL and OPENAI_EMBEDDING_MODEL and MILVUS_HOST and MILVUS_PORT):
        raise RuntimeError("缺少必要环境变量：OPENAI_API_KEY/OPENAI_BASE_URL/OPENAI_EMBEDDING_MODEL/MILVUS_HOST/MILVUS_PORT")

    emb = OpenAIEmbeddings(
        api_key=OPENAI_API_KEY,
        base_url=OPENAI_BASE_URL,
        model=OPENAI_EMBEDDING_MODEL,
    )
    q_emb = emb.embed_query(query.strip())

    client = MilvusClient(host=MILVUS_HOST, port=MILVUS_PORT)

    try:
        client.load_collection(COLLECTION_NAME)
    except Exception as e:
        logger.warning(f"加载集合 {COLLECTION_NAME} 失败（可能未创建或未插入）：{e}")

    results = client.search(
        collection_name=COLLECTION_NAME,
        data=[q_emb],
        limit=top_k,
        output_fields=["category", "title", "source", "content"],
    )

    hits = results[0] if results else []
    if not hits:
        print("未命中任何结果。请确认集合已创建且已插入数据。")
        return

    print(f"\nTop {top_k} 检索结果 for: {query}")
    for i, hit in enumerate(hits, 1):
        fields = hit.get("entity", {}) if isinstance(hit, dict) else hit
        # 兼容 pymilvus client 返回结构（取值方式可能因版本略有差异）
        title = fields.get("title") if isinstance(fields, dict) else getattr(fields, "title", "")
        source = fields.get("source") if isinstance(fields, dict) else getattr(fields, "source", "")
        content = fields.get("content") if isinstance(fields, dict) else getattr(fields, "content", "")
        score = hit.get("distance") if isinstance(hit, dict) else getattr(hit, "distance", None)

        snippet = content[:120].replace("\n", " ") + ("..." if len(content) > 120 else "")
        print(f"[{i}] score={score}  title={title}  source={source}\n    {snippet}")


def main():
    query = "报销多久到账" if len(sys.argv) < 2 else " ".join(sys.argv[1:])
    search(query)


if __name__ == "__main__":
    main()

